#!/usr/bin/env python3
# -*- coding: utf-8 -*-

from pathlib import Path
import re
import time
import shutil
import tempfile
import sys
import os
import difflib
import argparse
from contextlib import contextmanager
from typing import Optional, Tuple, Any

# ========= 常量 =========
DEFAULT_PATH = "/etc/pacman.conf"
DEFAULT_MIRROR = "https://mirrors.ustc.edu.cn/archlinuxcn/$arch"

SECTION_RE = re.compile(r"(?m)^\s*\[(?P<name>[^\]]+)\]\s*$")
OPT_LINE = lambda key: re.compile(rf"(?m)^(?P<indent>\s*){re.escape(key)}(?:\s*=\s*.*)?\s*$")
HAS_LINE = lambda key: re.compile(rf"(?m)^\s*{re.escape(key)}(?:\s*=\s*.*)?\s*$")
COMMENT_OR_BLANK = re.compile(r"(?m)^\s*(#.*)?$")


# ========= 纯文本变换（幂等，可单测） =========
def _span_of_section(text: str, name: str) -> Optional[Tuple[int, int]]:
    name = name.strip().lower()
    spots = list(SECTION_RE.finditer(text))
    for idx, m in enumerate(spots):
        if m.group("name").strip().lower() == name:
            start = m.end()
            end = spots[idx + 1].start() if idx + 1 < len(spots) else len(text)
            return start, end
    return None


def _ensure_in_section(text: str, section: str, line: str, unique_key: Optional[str] = None):
    span = _span_of_section(text, section)
    key_re = HAS_LINE(unique_key or line.strip().split()[0])
    if span:
        s, e = span
        before, body, after = text[:s], text[s:e], text[e:]
        if not key_re.search(body):
            indent = ""
            for ln in body.splitlines():
                if COMMENT_OR_BLANK.match(ln):
                    continue
                indent = re.match(r"\s*", ln).group(0)
                break
            if not body.endswith("\n"):
                body += "\n"
            body += f"{indent}{line.rstrip()}\n"
            return before + body + after, True
        return text, False
    else:
        if not text.endswith("\n"):
            text += "\n"
        block = f"\n[{section}]\n{line.rstrip()}\n"
        return text + block, True


def _replace_or_insert_siglevel(text: str, value: str = "Never"):
    # 仅在 [options] 段内替换/追加，避免误改其他仓库段
    span = _span_of_section(text, "options")
    if span:
        s, e = span
        before, body, after = text[:s], text[s:e], text[e:]
        sig_re = OPT_LINE("SigLevel")
        if sig_re.search(body):
            new_body = sig_re.sub(rf"\g<indent>SigLevel = {value}", body, count=1)
            changed = new_body != body
            return before + new_body + after, changed
        # [options] 存在但没有 SigLevel：追加
        if not body.endswith("\n"):
            body += "\n"
        body += f"SigLevel = {value}\n"
        return before + body + after, True
    # 没有 [options] 段：新建并写入
    return _ensure_in_section(text, "options", f"SigLevel = {value}", unique_key="SigLevel")


def _ensure_color(text: str):
    return _ensure_in_section(text, "options", "Color", unique_key="Color")


def _ensure_ilovecandy_below_color(text: str):
    if HAS_LINE("ILoveCandy").search(text):
        return text, False
    lines = text.splitlines(keepends=True)
    color_re = OPT_LINE("Color")
    for i, ln in enumerate(lines):
        m = color_re.match(ln)
        if m:
            indent = m.group("indent")
            lines.insert(i + 1, f"{indent}ILoveCandy\n")
            return "".join(lines), True
    return text, False


def _ensure_parallel_downloads(text: str, value: int = 5):
    """确保 [options] 里有 ParallelDownloads = <value>（去注释/替换值/若无则添加）"""
    pd_re = OPT_LINE("ParallelDownloads")
    span = _span_of_section(text, "options")
    if span:
        s, e = span
        before, body, after = text[:s], text[s:e], text[e:]
        if pd_re.search(body):
            new_body = pd_re.sub(rf"\g<indent>ParallelDownloads = {value}", body, count=1)
            return before + new_body + after, new_body != body
        if not body.endswith("\n"):
            body += "\n"
        body += f"ParallelDownloads = {value}\n"
        return before + body + after, True
    return _ensure_in_section(text, "options", f"ParallelDownloads = {value}", unique_key="ParallelDownloads")


def _ensure_repo(text: str, section="archlinuxcn", mirror=DEFAULT_MIRROR):
    # 段不存在：追加完整段
    sec_re = re.compile(rf"(?mi)^\s*\[{re.escape(section)}\]\s*$")
    m = sec_re.search(text)
    if not m:
        if not text.endswith("\n"):
            text += "\n"
        block = f"\n[{section}]\nServer = {mirror}\n"
        return text + block, True
    # 段已存在：确保有未注释的 Server 行
    s = m.end()
    next_m = next(sec_re.finditer(text, s), None)
    e = next_m.start() if next_m else len(text)
    before, body, after = text[:s], text[s:e], text[e:]
    has_server = re.search(r"(?mi)^\s*Server\s*=", body) is not None
    if has_server:
        return text, False
    if not body.endswith("\n"):
        body += "\n"
    body += f"Server = {mirror}\n"
    return before + body + after, True


def transform_text(
    content: str,
    *,
    mirror: str,
    siglevel: str,
    add_color: bool,
    add_candy: bool,
    parallel_downloads: Optional[int] = None,
):
    text = content
    changed = False

    text, c = _replace_or_insert_siglevel(text, siglevel)
    changed |= c

    if add_color:
        text, c = _ensure_color(text)
        changed |= c
    if add_candy:
        text, c = _ensure_ilovecandy_below_color(text)
        changed |= c
    if parallel_downloads is not None:
        text, c = _ensure_parallel_downloads(text, parallel_downloads)
        changed |= c

    text, c = _ensure_repo(text, section="archlinuxcn", mirror=mirror)
    changed |= c

    return text, changed


# ========= 备份/写入/锁 =========
def resolve_backup_dir(target: Path, user_given: Optional[str]) -> Path:
    if user_given:
        return Path(user_given).expanduser().resolve()
    return target.parent  # 先试同目录


def fallback_state_dir() -> Path:
    xdg = os.environ.get("XDG_STATE_HOME")
    base = Path(xdg) if xdg else Path.home() / ".local" / "state"
    return base / "archlinuxcn-conf" / "backups"


def backup_file(src: Path, backup_dir: Path) -> Path:
    backup_dir = backup_dir.expanduser().resolve()
    ts = time.strftime("%Y%m%d-%H%M%S")
    backup_dir.mkdir(parents=True, exist_ok=True)
    dst = backup_dir / (src.name + f".bak.{ts}")
    try:
        shutil.copy2(src, dst)
        return dst
    except PermissionError:
        fb = fallback_state_dir()
        fb.mkdir(parents=True, exist_ok=True)
        dst = fb / (src.name + f".bak.{ts}")
        shutil.copy2(src, dst)
        return dst


@contextmanager
def locked_file(path: Path, mode="r+", encoding="utf-8"):
    path.touch(exist_ok=True)
    f = path.open(mode, encoding=encoding)
    try:
        if os.name == "posix":
            try:
                import fcntl
                fcntl.flock(f.fileno(), fcntl.LOCK_EX)
            except Exception:
                pass
        yield f
    finally:
        try:
            f.flush()
        except Exception:
            pass
        if os.name == "posix":
            try:
                import fcntl
                fcntl.flock(f.fileno(), fcntl.LOCK_UN)
            except Exception:
                pass
        f.close()


def atomic_write(path: Path, content: str, final_mode: int = 0o644, preserve_owner: bool = True):
    """
    原子写入：确保最终权限为 final_mode（默认 0644），并尽可能保留原文件属主/属组。
    """
    # 尝试获取原文件的 stat（存在且可读时）
    st = None
    try:
        st = os.stat(path)
    except Exception:
        st = None

    tmp_fd, tmp_path = tempfile.mkstemp(prefix=".pacman_conf.", dir=str(path.parent))
    try:
        with os.fdopen(tmp_fd, "w", encoding="utf-8", newline="") as tmp:
            tmp.write(content)
            tmp.flush()
            os.fsync(tmp.fileno())

        # 先给临时文件设定目标权限
        try:
            os.chmod(tmp_path, final_mode)
        except Exception as e:
            print(f"[!] 设置临时文件权限失败（继续执行）：{e}", file=sys.stderr)

        # 若需要保留属主/属组，则在替换前先对临时文件设置一致的属主/属组（需要权限）
        if preserve_owner and st is not None:
            try:
                os.chown(tmp_path, st.st_uid, st.st_gid)
            except Exception:
                # 没权限也没关系，替换后再尝试
                pass

        # 原子替换
        os.replace(tmp_path, path)

        # 替换后再兜底：权限与属主/属组
        try:
            os.chmod(path, final_mode)
        except Exception as e:
            print(f"[!] 设置最终文件权限失败：{e}", file=sys.stderr)

        if preserve_owner and st is not None:
            try:
                os.chown(path, st.st_uid, st.st_gid)
            except Exception:
                # 非特权用户修改 /etc/* 可能失败，这里容错
                pass
    finally:
        # 若 replace 失败或异常中断，清理残留临时文件
        try:
            if os.path.exists(tmp_path):
                os.unlink(tmp_path)
        except Exception:
            pass


# ========= 主流程 =========
def archlinuxcn_source(
    path: str = DEFAULT_PATH,
    *,
    mirror: str = DEFAULT_MIRROR,
    siglevel: str = "Never",
    add_color: bool = True,
    add_candy: bool = True,
    parallel_downloads: Optional[int] = None,
    backup_dir: Optional[str] = None,
    do_backup: bool = True,
    dry_run: bool = False,
    dry_run_format: str = "diff",  # diff/context/summary
    require_root: bool = False,
) -> int:
    p = Path(path)

    if not p.exists():
        print(f"[!] 文件不存在：{p}", file=sys.stderr)
        return 2

    # 系统路径：非 dry-run 时给出权限/身份提示
    sys_path = str(p).startswith("/etc/")
    if not dry_run and sys_path:
        if require_root and hasattr(os, "geteuid") and os.geteuid() != 0:
            print(f"[!] 要修改 {p} 需要 root。请使用 sudo，或先 -n 预览、-p 指向可写副本测试。", file=sys.stderr)
            return 2
        if not os.access(p, os.W_OK) and not os.access(p.parent, os.W_OK):
            print(f"[!] 对 {p} 无写权限。请使用 sudo，或 -n 预览、-p 指向可写副本测试。", file=sys.stderr)
            return 2

    original = p.read_text(encoding="utf-8", errors="ignore")
    new_text, changed = transform_text(
        original,
        mirror=mirror,
        siglevel=siglevel,
        add_color=add_color,
        add_candy=add_candy,
        parallel_downloads=parallel_downloads,
    )

    if dry_run:
        if changed:
            if dry_run_format == "diff":
                print("=== DRY-RUN (unified diff) ===")
                diff = difflib.unified_diff(
                    original.splitlines(keepends=True),
                    new_text.splitlines(keepends=True),
                    fromfile=f"{p} (before)",
                    tofile=f"{p} (after)",
                )
                sys.stdout.writelines(diff)
            elif dry_run_format == "context":
                print("=== DRY-RUN (context diff) ===")
                diff = difflib.context_diff(
                    original.splitlines(keepends=True),
                    new_text.splitlines(keepends=True),
                    fromfile=f"{p} (before)",
                    tofile=f"{p} (after)",
                )
                sys.stdout.writelines(diff)
            elif dry_run_format == "summary":
                print("=== DRY-RUN (summary) ===")
                if re.search(r"(?m)^\s*SigLevel\s*=", original):
                    if not re.search(r"(?m)^\s*SigLevel\s*=\s*Never\s*$", original) and \
                       re.search(r"(?m)^\s*SigLevel\s*=\s*Never\s*$", new_text):
                        print("  * 更新 SigLevel -> Never")
                else:
                    print("  * 添加 SigLevel = Never")
                if not re.search(r"(?m)^\s*Color\s*$", original) and re.search(r"(?m)^\s*Color\s*$", new_text):
                    print("  * 添加 Color")
                if "ILoveCandy" not in original and "ILoveCandy" in new_text:
                    print("  * 添加 ILoveCandy")
                if parallel_downloads is not None:
                    if re.search(r"(?m)^\s*ParallelDownloads\s*=", original):
                        print(f"  * 更新 ParallelDownloads -> {parallel_downloads}")
                    else:
                        print(f"  * 添加 ParallelDownloads = {parallel_downloads}")
                if "[archlinuxcn]" not in original and "[archlinuxcn]" in new_text:
                    print("  * 添加 [archlinuxcn] 段/Server")
        else:
            print("DRY-RUN: 无需修改（已符合目标状态）。")
        return 0

    if not changed:
        print("无需修改：文件已是目标状态。")
        return 0

    if do_backup:
        try:
            dst = backup_file(p, resolve_backup_dir(p, backup_dir))
            print(f"已备份到：{dst}")
        except Exception as e:
            print(f"[!] 备份失败（继续写入）：{e}", file=sys.stderr)

    with locked_file(p, mode="r+", encoding="utf-8") as f:
        cur = f.read()
        if cur != original:
            new_text, _ = transform_text(
                cur,
                mirror=mirror,
                siglevel=siglevel,
                add_color=add_color,
                add_candy=add_candy,
                parallel_downloads=parallel_downloads,
            )
        atomic_write(p, new_text, final_mode=0o644, preserve_owner=True)

    print("写入完成。")
    return 0


# ========= CLI =========
def parse_args():
    ap = argparse.ArgumentParser(
        description="幂等修改 pacman.conf：SigLevel/Color/ILoveCandy/ParallelDownloads/archlinuxcn 段。"
    )
    ap.add_argument("-p", "--path", default=DEFAULT_PATH,
                    help=f"目标配置文件路径（默认 {DEFAULT_PATH}）")
    ap.add_argument("-m", "--mirror", default=DEFAULT_MIRROR,
                    help="archlinuxcn 镜像地址")
    ap.add_argument("-s", "--siglevel", default="Never",
                    help="SigLevel 值（默认 Never）")
    ap.add_argument("-P", "--parallel-downloads", type=int, default=5,
                    help="设置 ParallelDownloads 的值（例如 5/10），默认不改动")
    ap.add_argument("-C", "--no-color", action="store_true",
                    help="不添加 Color")
    ap.add_argument("-c", "--no-candy", action="store_true",
                    help="不添加 ILoveCandy")
    ap.add_argument("-B", "--no-backup", action="store_true",
                    help="不备份原文件")
    ap.add_argument("-d", "--backup-dir", default=None,
                    help="备份目录（无写权限时回退到 XDG_STATE_HOME 或 ~/.local/state/...）")
    ap.add_argument("-n", "--dry-run", action="store_true",
                    help="仅预览修改，不落盘")
    ap.add_argument("-f", "--format", dest="dry_run_format",
                    choices=["diff", "context", "summary"], default="summary",
                    help="dry-run 输出样式：diff/context/summary（默认 summary）")
    ap.add_argument("-r", "--require-root", action="store_true",
                    help="在系统路径上需要 root 身份（更安全）")
    return ap.parse_args()


def main():
    args = parse_args()
    rc = archlinuxcn_source(
        path=args.path,
        mirror=args.mirror,
        siglevel=args.siglevel,
        add_color=not args.no_color,
        add_candy=not args.no_candy,
        parallel_downloads=args.parallel_downloads,
        backup_dir=args.backup_dir,
        do_backup=not args.no_backup,
        dry_run=args.dry_run,
        dry_run_format=args.dry_run_format,
        require_root=args.require_root,
    )
    sys.exit(rc)


if __name__ == "__main__":
    main()
